#include "polygon.h"

//#define VLADWINDOW 1 /* Erm...  *shrugs* */
//#define PRINT_CONTOURS 1
//#define PRINT_POLYGON 1
//#define PRINT_RECT 1
#define INTERACTIVE 1

// #define DEBUG_RATIO

#define ABS(x) (x > 0) ? x : -x

//const int MIN_AREA = 2600;
const int MIN_AREA = 4500;
const double MIN_AREA_RATIO = 0.7;
const int MAX_LENGTH_DIFFERENCE = 400;

// enable testing with small hand held target
//#define TESTING_SMALL_TARGET

#ifndef TESTING_SMALL_TARGET
#define TOP_RECT_RATIO_MIN .25
#define TOP_RECT_RATIO_MAX .41
#define MIDDLE_RECT_RATIO_MIN .41
#define MIDDLE_RECT_RATIO_MAX .54
#else
#define TOP_RECT_RATIO_MIN .7
#define TOP_RECT_RATIO_MAX 1.0
#define MIDDLE_RECT_RATIO_MIN 1.0
#define MIDDLE_RECT_RATIO_MAX 1.0
#endif

using namespace cv;
using namespace std;

void saveFrame(char *fileBase, Mat *frame)
{
    size_t fbsize, cursize;
    char buf[80];
    struct tm *tmp;
    struct timeval tv; 
    int milliseconds; 

 
    /* Obtain the time of day, and convert it to a tm struct. */ 
    gettimeofday (&tv, NULL); 

    tmp = localtime (&tv.tv_sec);
    milliseconds = tv.tv_usec / 1000;

    fbsize=strlen(fileBase);
    
    strcpy(buf, fileBase);
    
    cursize = strftime(&buf[fbsize], 80 - fbsize, "%T", tmp);
    sprintf(&buf[fbsize + cursize], ".%3.3d.jpg", milliseconds);
#ifdef DEBUG
    printf("wrote: %s\n", buf);
#endif
    imwrite(buf, *frame);
    
}

RectangleScanner::RectangleScanner() {}

void RectangleScanner::SetImage(Mat image){
    src = image.clone();
    src_copy = image.clone();
#if defined(SAVE_IMAGES) || defined(WINDOW)
    src_rects = image.clone();
#endif /* SAVE_IMAGES */
    rectangles.clear();
}

void RectangleScanner::PrepareImage(){

#ifdef WINDOW
    imshow("original", src);
    waitKey(0);
#endif /* WINDOW */
    
    //Preform a color threshold
    //ColorThreshold(src_copy);
    FilterColor(src);

#ifdef WINDOW
    imshow("Test1: Filter Color", src);
    waitKey(0);
#endif /* WINDOW */

    //Convert thresholded image to grayscale
    cvtColor(src, gray, CV_BGR2GRAY);

#ifdef WINDOW
    imshow("test2: Grayscale", gray);
    waitKey(0);
#endif /* WINDOW */
    

    //Threshold the grayscale image
    threshold(gray, gray, 190, 255, THRESH_BINARY_INV);

#if WINDOW
    imshow("test3: Binary threshold", gray);
    waitKey(0);
#endif /* WINDOW */

    //find contours
    vector<Vec4i> prelim_hierarchy;
    vector<vector<Point> > prelim_contours;
    findContours(gray, prelim_contours, prelim_hierarchy, CV_RETR_LIST, CV_CHAIN_APPROX_NONE);

#ifdef DEBUG    //debug output
    cout << "prelim_contours.size() = " << prelim_contours.size() << endl;
#endif

    //redraw contours
    gray = Scalar::all(255);
    for (size_t i=0; i<prelim_contours.size(); i++)
    {
        drawContours(gray, prelim_contours, i, Scalar(0), -1);
    }

#ifdef WINDOW                                                                   
    imshow("test4: Created and drew contours", gray);
    waitKey(0);
#endif /* WINDOW */


    //Kernel we use to later dilate and erote image
    Mat element = getStructuringElement(MORPH_RECT, Size(2, 2), Point(1, 1));
    dilate(gray, gray, element);

#ifdef WINDOW
    imshow("test5: Dilate", gray);
    waitKey(0);
#endif /* WINDOW */

    //Erode image
    erode(gray, gray, element);
    
#ifdef WINDOW
    imshow("test6: Erode", gray);
    waitKey(0);
#endif /* WINDOW */
}


void RectangleScanner::FindContours(){

    //Prepare the image with a grayscale/dilation/erosion first
    PrepareImage();

    //Find contours
    findContours(gray, contours, hierarchy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);

#ifdef DEBUG
    cout << "contours.size() = " << contours.size() << endl;
#endif

} 

void RectangleScanner::FindRectangles(){

    //First, find all the contours
    FindContours();
   
    //variables for processing 
    unsigned size = contours.size();
    bool shape_is_rectangle = false;

    for (size_t i = 0; i < (size); ++i){

#ifdef DEBUG
        cout << endl << "contours[" << i << "]:" << endl;
#endif
        //access new element 
        vector<Point> polygon = contours[i];

#ifdef PRINT_CONTOURS
        cout << "num points = " << polygon.size() << endl;
        for (size_t j=0; j< polygon.size(); j++) {
            cout << "point["<<j<<"] = " << polygon[j] << endl;
        }
#endif /* PRINT_CONTOURS */

        vector<Point> coordinates;
        /* Need to turn epsilon value */
        approxPolyDP(polygon, coordinates, 5, true);

#ifdef PRINT_POLYGON
        cout << "num points = " << coordinates.size() << endl;
        for (size_t j=0; j< coordinates.size(); j++) {
            cout << "coordinates["<<j<<"] = " << coordinates[j] << endl;
        }
#endif /* PRINT_POLYGON */

        /* minAreaRect: note it is rotated rectangle */
        RotatedRect rrect = minAreaRect(polygon);
        double polarea = contourArea(polygon);
        double rectarea = rrect.size.height*rrect.size.width;

#ifdef DEBUG
        cout << "Contour area is " << polarea << endl;
        cout << "Rect area is " << rectarea <<  endl;
#endif

        //If the shape is close in area to its smallest-area rectangle, it's a rectangle
        //If it looks like a rectangle and walks like a rectangle, it's probably a rectangle
        shape_is_rectangle = false;
        if (polarea/rectarea > MIN_AREA_RATIO) {
            shape_is_rectangle = true;

#ifdef DEBUG
            cout << "That's a rectangle, ratio is only " << polarea/rectarea << endl;
#endif

        }
        if (shape_is_rectangle){
            rectangles.push_back(rrect);
        }

        //Reset the bool after usage
        shape_is_rectangle = false;

#ifdef PRINT_RECT        
        cout << "Min Area Rectangle:" << " center = " << rrect.center;
        cout << " height = " << rrect.size.height;
        cout << " width = " << rrect.size.width;
        cout << " angle = " << rrect.angle << endl;
#endif
    }

}

void RectangleScanner::ChooseRectangle(){

    FindRectangles();

    unsigned size = rectangles.size();
    RotatedRect rrect;

    bool shape_is_large = true;
    bool shape_is_image = true;
    bool shape_is_on_edge = true;
    
    top_rect_selected = false;
    middle_rect_selected = false;
    RECT current_rect;

#ifdef DEBUG
    cout << "rectangles.size() = " << size << endl;
#endif

    for (size_t i = 0; i < (size); ++i) {

        rrect = rectangles[i];

        shape_is_large = false;
        shape_is_image = false;
        shape_is_on_edge = true;

        middle_rect_selected = false;

        //Check that shape fulfills minimum size requirement
        if (rrect.size.height*rrect.size.width > MIN_AREA) shape_is_large = true;

        //Check if the shape is the entire image (algorithm sometimes finds it)
        if ( rrect.size.height * rrect.size.width > 290000) {
            shape_is_image = true;
        }

        //Find the real dimensions of the rectangle
        double height = rrect.size.height;
        double width  = rrect.size.width;
#if 0 // better use direct height vs width comparion rather than angle assumption
        if (rrect.angle > 45 || rrect.angle < -45) {
            height = rrect.size.width;
            width  = rrect.size.height;
            rrect.size.height = height;
            rrect.size.width = width;
        }
#else
        if (height > width) {
            height = rrect.size.width;
            width  = rrect.size.height;
            
            rrect.size.height = height;
            rrect.size.width = width;            
        }
#endif /* 0 */

        if (rrect.angle > 45) {
            rrect.angle = 90 - rrect.angle;
        } else if (rrect.angle < -45) {
            rrect.angle = 90 + rrect.angle;
        } else {}

        
        if (TOP_RECT_RATIO_MIN < height/width && height/width < TOP_RECT_RATIO_MAX) {
            current_rect = TOP;
        } else if  (.41 < height/width && height/width < .54) {
            current_rect = MIDDLE;
            middle_rect_selected = true;
        } else {
            current_rect = OTHER; 
        }

        if (((rrect.center.x - (width / 2)) > 10) &&
            ((rrect.center.x + (width / 2)) < 630)) {
             shape_is_on_edge = false;
        }

#if defined(SAVE_IMAGES) || defined(WINDOW)
        // let's draw all big rectangles that we detected       
        if (shape_is_large && !(shape_is_image)) {
#ifdef DEBUG_RATIO
            cout << "rrect.center.x = " << rrect.center.x << endl;
            cout << "rrect.center.y = " << rrect.center.y << endl;
            cout << "rrect.angle = " <<  rrect.angle  << endl;
            cout << "height = " << height << endl;
            cout << "width = " << width << endl;
            cout << "height/width = " << height/width << endl;
            cout << "height * width = " << height * width << endl;
            cout << endl;
#endif /* DEBUG_RATIO */
            Point2f vertices[4];
            rrect.points(vertices);
            for (int j = 0; j < 4; j++){
                line(src_rects, vertices[j], vertices[(j+1)%4], Scalar(0,0,255));
            }
            
            circle(src_rects, rrect.center, 10, Scalar(255,0,0));
        }
#endif /* SAVE_IMAGES */
       
        //fitness rating (checking various booleans)
        if (shape_is_large && !(shape_is_image) && (!shape_is_on_edge)) {

            if (current_rect == TOP){
                best_match_rect = rrect;
                top_rect_selected = true;

#if defined(SAVE_IMAGES) || defined(WINDOW)
                // mark one we selected with solid color
                circle(src_rects, best_match_rect.center, 10, Scalar(255,0,0), -1);
#endif /* SAVE_IMAGES */
            }
#ifdef DEBUG
            cout << i + 1 << " - This one works" << endl;
#endif
        }
    }


    
#if WINDOW
    imshow("rectangle", src_rects);
    waitKey(0);
    destroyAllWindows();
#endif
}

RotatedRect RectangleScanner::GetRectangle(){
    return best_match_rect;
}

